from ryu.controller import ofp_event
from ryu.controller.handler import MAIN_DISPATCHER, DEAD_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.lib import hub
import simple_switch_stp_13

from datetime import datetime

import sys
sys.path.append('..')

import os
if not os.path.exists('../data'): os.makedirs('../data')

from classifier.information_entropy import Entropy


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

class CollectTrainingStatsApp(simple_switch_stp_13.SimpleSwitch13):
    def __init__(self, *args, **kwargs):
        super(CollectTrainingStatsApp, self).__init__(*args, **kwargs)
        self.datapaths = {}
        self.monitor_thread = hub.spawn(self.monitor)
        self.entropy = Entropy()    # 香农信息熵
        self.flow_type = 1    # 收集的流量类别   0:正常流量    1:攻击流量

    # Asynchronous message
    @set_ev_cls(ofp_event.EventOFPStateChange, [MAIN_DISPATCHER, DEAD_DISPATCHER])
    def state_change_handler(self, ev):
        datapath = ev.datapath
        if ev.state == MAIN_DISPATCHER:
            if datapath.id not in self.datapaths:
                self.logger.debug('register datapath: %016x', datapath.id)
                self.datapaths[datapath.id] = datapath

        elif ev.state == DEAD_DISPATCHER:
            if datapath.id in self.datapaths:
                self.logger.debug('unregister datapath: %016x', datapath.id)
                del self.datapaths[datapath.id]

    def monitor(self):
        while True:
            for dp in self.datapaths.values():
                self.request_stats(dp)
            hub.sleep(10)

    def request_stats(self, datapath):
        self.logger.debug('######## send stats request: %016x', datapath.id)

        parser = datapath.ofproto_parser

        req = parser.OFPFlowStatsRequest(datapath)
        datapath.send_msg(req)

    @set_ev_cls(ofp_event.EventOFPFlowStatsReply, MAIN_DISPATCHER)
    def _flow_stats_reply_handler(self, ev):
        # 写入文件
        with open('../data/flow_stat.csv', 'a+', encoding='utf-8') as f:
            body = ev.msg.body
            for stat in sorted([flow for flow in body if (flow.priority == 1)],
                               key=lambda flow:(
                                   flow.match['eth_type'],
                                   flow.match['ipv4_src'],
                                   flow.match['ipv4_dst'],
                                   flow.match['ip_proto']
                               )):

                ip_src = stat.match['ipv4_src']
                ip_dst = stat.match['ipv4_dst']

                ip_proto = stat.match['ip_proto']    # column 1
                flow_duration_sec = stat.duration_sec    # column 2
                idle_timeout = stat.idle_timeout    # column 3
                hard_timeout = stat.hard_timeout    # column 4
                packet_count = stat.packet_count    # fcolumn 5
                # column 6
                try:
                    packet_count_per_second = stat.packet_count / stat.duration_sec
                except:
                    packet_count_per_second = 0

                byte_count = stat.byte_count    # column 7
                # column 8
                try:
                    byte_count_per_second = stat.byte_count / stat.duration_sec
                except:
                    byte_count_per_second = 0

                self.entropy.collect(ip_dst)
                entro_val = self.entropy.value    # column 9

                f.write('{},{},{},{},{},{},{},{},{},{}\n'.format(
                    ip_proto, flow_duration_sec, idle_timeout, hard_timeout, packet_count,
                    packet_count_per_second, byte_count, byte_count_per_second, entro_val, self.flow_type
                ))

